#include "LearningAndRecognition.h"
#include "FacialRecognition.h"
#include <stdio.h>
#include <string.h>

const char * TRAIN_FACES = "train_faces.txt";
const char * RECOGNIZE_FACES = "recognize_face.txt";

//// Global variables
int nTrainFaces = 0; // number of training images
int nEigens = 0; // number of eigenvalues
IplImage ** faceImgArr = 0; // array of face images
CvMat * personNumTruthMat = 0; //array of person numbers
IplImage * pAvgTrainImg = 0; // the average image
IplImage ** eigenVectArr = 0; // eigenvectors
CvMat * eigenValMat = 0; // eigenvalues
CvMat * projectedTrainFaceMat = 0; // projected training faces
double leastDistSq = DBL_MAX; //Least distance to image
const double RECOGNITION_THRESHOLD = 0.20;
int recognitionCounter = 0;

LearningAndRecognition::LearningAndRecognition(void)
{
}

void LearningAndRecognition::LearnAndRecognizeFace(char * argv[])
{
	std::fstream file;
	std::string optionInput = "";
	std::string input = "";
	std::vector<std::string> vectorFilePaths;
	std::vector<std::string> vectorPersonNames;
	int personCounter = 0;

	std::string filePath = "";
	int counter = 0;
	file.open(TRAIN_FACES);
	while(!file.eof()) {
		getline(file, filePath);
		if(filePath.empty())
			break;
		counter++;
	}
	file.close();
	std::cout << "\n1. Recognize my face.\n";
	//getline(std::cin, optionInput);

	//Need at least 3 faces to start recognition phase
	if(counter >= 3) {

		/*std::stringstream ss;
		std::string facesDataString;
		ss << argv[1];
		ss >> facesDataString;
		facesDataString += "\\facesdata.xml";
		const char * FACES_DATA = facesDataString.c_str();*/

		//Try to recognize face only if there is a DB of faces available
		if(fexists("facesdata.xml")) {
			while(leastDistSq >= RECOGNITION_THRESHOLD) {
				recognize(argv, vectorFilePaths);
				recognitionCounter++;

				if(recognitionCounter == 10) {
					std::cout << "\nYour face could not be recognized! Please train your face and try again...\n";
					recognitionCounter = 0;
					break;
				}
			}
			leastDistSq = DBL_MAX;
		}
		else
			std::cout << "\n" << "facesdata.xml" << "is missing! Please train faces again...\n";
	}

	else {
		std::cout << "\nAt least 3 faces must be learned first! Press any key to exit...";
	}
}

std::wstring LearningAndRecognition::s2ws(const std::string& s)
{
	int len;
	int slength = (int)s.length() + 1;
	len = MultiByteToWideChar(CP_ACP, 0, s.c_str(), slength, 0, 0); 
	wchar_t* buf = new wchar_t[len];
	MultiByteToWideChar(CP_ACP, 0, s.c_str(), slength, buf, len);
	std::wstring r(buf);
	delete[] buf;
	return r;
}

LearningAndRecognition::~LearningAndRecognition(void)
{
}

void LearningAndRecognition::CreateRecognitionTextFileWithFilePaths(std::vector<std::string> &vectorFilePaths)
{
	std::ofstream file;
	file.open(RECOGNIZE_FACES);
	file << vectorFilePaths[0] + "\n";
	file.close();
	vectorFilePaths.clear();
}

void LearningAndRecognition::ReadTrainingTextFile(std::map<int, std::string> &mapPersonNumName)
{
	std::fstream file;
	std::string filePath = "";
	int personNum  = 0;
	std::string name = "";
	file.open(TRAIN_FACES);

	while(!file.eof()) {
		getline(file, filePath);
		if(filePath.empty() && mapPersonNumName.empty()) {
			std::cout << "\nTraining text file is empty! Exiting...\n";
			exit(1);			
			std::terminate();
		}
		else {
			if(filePath.empty())
				break;
			//Store the person number and name
			int i = 0;
			std::string tempNum = "";
			for(i; i < filePath.length(); i++) {
				if(filePath[i] == ' ') {
					i++;
					break;
				}
				tempNum += filePath[i];				
			}
			personNum = atoi(tempNum.c_str());
			tempNum = "";
			for(i; i < filePath.length(); i++) {
				if(filePath[i] == '\\')
					break;
				name += filePath[i];
			}
			mapPersonNumName[personNum] = name;
			name = "";
		}		
	}
	
	file.close();
}

bool LearningAndRecognition::fexists(const char * FILE)
{
	std::ifstream file(FILE);
	return file;
}

int LearningAndRecognition::loadFaceImgArray(const char * filename)
{
	FILE * imgListFile = 0;
	char imgFilename[512];
	int iFace, nFaces=0;

	// open the input file
	imgListFile = fopen(filename, "r");

	// count the number of faces
	while( fgets(imgFilename, 512, imgListFile) ) ++nFaces;
	rewind(imgListFile);

	// allocate the face-image array and person number matrix
	faceImgArr = (IplImage **)cvAlloc( nFaces*sizeof(IplImage *) );
	personNumTruthMat = cvCreateMat( 1, nFaces, CV_32SC1 );

	// store the face images in an array
	for(iFace=0; iFace<nFaces; iFace++)
	{
		// read person number and name of image file
		fscanf(imgListFile,
			"%d %s", personNumTruthMat->data.i+iFace, imgFilename);

		// load the face image
		faceImgArr[iFace] = cvLoadImage(imgFilename, CV_LOAD_IMAGE_GRAYSCALE);
	}

	fclose(imgListFile);

	return nFaces;
}

void LearningAndRecognition::recognize(char * argv[], std::vector<std::string> vectorFilePaths)
{
	CapturingAndStoring capturingAndStoring;
	std::map<int, std::string> mapPersonNumName;
	int i, nTestFaces = 0; // the number of test images
	CvMat * trainPersonNumMat = 0; // the person numbers during training
	float * projectedTestFace = 0;
	std::vector<std::string> vectorFilePath;
	std::string currentInput = "Temp";
	std::wstring stemp = s2ws(currentInput); //Temporary buffer is required
	LPCWSTR dir = stemp.c_str();
	CreateDirectory(dir, NULL);

	//Prepare the current face image for capturingAndStoring
	capturingAndStoring.PrepareFaceInImageAndStoreOnDisk(argv, FacialRecognition::inputImage, FacialRecognition::faceRect, currentInput, 1, "", "", false, vectorFilePath);

	if(CapturingAndStoring::killThread) {
		std::cout << "\nThe square has to be green! Please exit and try again...\n";
		exit(1);			
		std::terminate();
	}

	//Create text file for the taken face image
	CreateRecognitionTextFileWithFilePaths(vectorFilePath);

	//Read training text file to grab person number and name for output
	ReadTrainingTextFile(mapPersonNumName);

	// load test images and ground truth for person number
	nTestFaces = loadFaceImgArray(RECOGNIZE_FACES);
	//printf("%d test faces loaded\n", nTestFaces);

	// load the saved training data
	if( !loadTrainingData( &trainPersonNumMat ) ) return;

	// project the test images onto the PCA subspace
	projectedTestFace = (float *)cvAlloc( nEigens*sizeof(float) );
	for(i=0; i<nTestFaces; i++)
	{
		int iNearest, nearest, truth;

		//project the test image onto the PCA subspace
		cvEigenDecomposite(
			faceImgArr[i],
			nEigens,
			eigenVectArr,
			0, 0,
			pAvgTrainImg,
			projectedTestFace);

		iNearest = findNearestNeighbor(projectedTestFace);
		truth = personNumTruthMat->data.i[i];
		nearest = trainPersonNumMat->data.i[iNearest];

		if(leastDistSq <= RECOGNITION_THRESHOLD)
			//printf("nearest = %d, Truth = %d\n", nearest, truth);
			std::cout << "\nYou are: " + mapPersonNumName[nearest] << "\n";
	}
}

int LearningAndRecognition::loadTrainingData(CvMat ** pTrainPersonNumMat)
{
	CvFileStorage * fileStorage;
	int i;

	// create a file-storage interface
	fileStorage = cvOpenFileStorage( "facesdata.xml", 0, CV_STORAGE_READ );
	if( !fileStorage )
	{
		fprintf(stderr, "Can't open facesdata.xml\n");
		return 0;
	}

	nEigens = cvReadIntByName(fileStorage, 0, "nEigens", 0);
	nTrainFaces = cvReadIntByName(fileStorage, 0, "nTrainFaces", 0);
	*pTrainPersonNumMat = (CvMat *)cvReadByName(fileStorage, 0, "trainPersonNumMat", 0);
	eigenValMat = (CvMat *)cvReadByName(fileStorage, 0, "eigenValMat", 0);
	projectedTrainFaceMat =
		(CvMat *)cvReadByName(fileStorage, 0, "projectedTrainFaceMat", 0);
	pAvgTrainImg = (IplImage *)cvReadByName(fileStorage, 0, "avgTrainImg", 0);
	eigenVectArr = (IplImage **)cvAlloc(nTrainFaces*sizeof(IplImage *));
	for(i=0; i<nEigens; i++)
	{
		char varname[200];
		sprintf( varname, "eigenVect_%d", i );
		eigenVectArr[i] = (IplImage *)cvReadByName(fileStorage, 0, varname, 0);
	}

	// release the file-storage interface
	cvReleaseFileStorage( &fileStorage );

	return 1;
}

int LearningAndRecognition::findNearestNeighbor(float * projectedTestFace)
{
	int i, iTrain, iNearest = 0;

	for(iTrain=0; iTrain<nTrainFaces; iTrain++)
	{
		double distSq=0;

		for(i=0; i<nEigens; i++)
		{
			float d_i =
				projectedTestFace[i] -
				projectedTrainFaceMat->data.fl[iTrain*nEigens + i];
			distSq += d_i*d_i/eigenValMat->data.fl[i];
		}

		if(distSq < leastDistSq)
		{
			leastDistSq = distSq;
			iNearest = iTrain;
		}
	}

	if(leastDistSq <= RECOGNITION_THRESHOLD)
		std::cout << "\nLeast distance: " << leastDistSq << "\n";
		//printf("%f", leastDistSq);
	return iNearest;
}
